import os
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torch.autograd import Variable
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
import datetime


class ConvBlock(nn.Module):
    def __init__(self):
        super(ConvBlock, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, stride=1, padding=1),  # 3*60*160 → 32*60*160
            nn.BatchNorm2d(32),  # 批归一化
            nn.LeakyReLU(0.2, inplace=True),  # 允许小于0的值通过（设置负斜率0.2），避免ReLU的“死亡神经元”问题。
            nn.MaxPool2d(kernel_size=2),  # 最大池化  32*60*160 → 32*30*80
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),  # 32*30*80 → 64*30*80
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2, inplace=True),
            nn.MaxPool2d(kernel_size=2),  # 64*30*80 → 64*15*40
            nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, stride=1, padding=1),  # 64*15*40 → 64*15*40
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2, inplace=True),
            nn.MaxPool2d(kernel_size=2)  # 64*15*40 → 64*7*20
        )
        self.fc1 = nn.Linear(in_features=64 * 7 * 20, out_features=512, )
        self.fc2 = nn.Linear(in_features=512, out_features=40)

    def forward(self, x):
        # 使用卷积提取特征
        x = self.conv(x)

        # 将特征图拉伸
        x = x.view(x.size(0), -1)

        # 使用输出层进行分类
        x = self.fc1(x)
        x = F.leaky_relu(x)
        x = self.fc2(x)
        return x

    def loss_function(self, output, label):
        loss = nn.CrossEntropyLoss()

        # 将验证码中每一个数字作为一个独立样本，然后分类
        output = output.contiguous().view(-1, 10)
        label = label.contiguous().view(-1)

        total_loss = loss(output, label)
        return total_loss

    def open_log_file(self, log_file_path=None):
        if os.path.exists(log_file_path):
            os.remove(log_file_path)
        log_file = open(log_file_path, 'a', encoding='utf-8')
        return log_file

    def close_log_file(self, log_file=None):
        log_file.close()

    def log(self, msg='', print_msg=True, end='\n', log_file=None):
        # 控制台打印信息
        if print_msg:
            print(msg, end=end)

        # 获取当前时间
        now = datetime.datetime.now()
        t = datetime.datetime.strftime(now, '%Y-%m-%d %H:%M:%S')

        # 写入日志文件
        if isinstance(log_file, str):
            log_file = self.open_log_file(log_file)
        if log_file is not None:
            log_file.write(t + ' ' + msg + '\n')
            log_file.flush()
